
#define INSTR_NEAR_PREFIX 0x0F
#define INSTR_SHORTJCC_BEGIN 0x70
#define INSTR_SHORTJCC_END 0x7F
#define INSTR_NEARJCC_BEGIN 0x80	//	Near's are prefixed with a 0x0F byte
#define INSTR_NEARJCC_END 0x8F
#define INSTR_RET 0xC2
#define INSTR_RETN 0xC3
#define INSTR_RETFN 0xCA
#define INSTR_RETF 0xCB
#define INSTR_RELJCX 0xE3
#define INSTR_RELJMP 0xE9
#define INSTR_SHORTJMP 0xEB

#define O_UNIQUE 0
#define O_PREFIX 1
#define O_IMM8   2
#define O_IMM16  3
#define O_IMM24  4
#define O_IMM32  5
#define O_IMM48  6
#define O_MODRM  7
#define O_MODRM8 8
#define O_MODRM32 9
#define O_EXTENDED 10
#define O_WEIRD   11
#define O_ERROR   12

__declspec(naked)int __cdecl mlde32(__in void*codeptr)
{
	__asm
	{
	   pushad

       cld
       xor  edx, edx

       mov  esi, [esp+(8*4)+4]
       mov  ebp, esp

       // 256 bytes, index-compressed opcode type table
       push 01097F71Ch
       push 0F71C6780h
       push 017389718h
       push 0101CB718h
       push 017302C17h
       push 018173017h
       push 0F715F547h
       push 04C103748h
       push 0272CE7F7h
       push 0F7AC6087h
       push 01C121C52h
       push 07C10871Ch
       push 0201C701Ch
       push 04767602Bh
       push 020211011h
       push 040121625h
       push 082872022h
       push 047201220h
       push 013101419h
       push 018271013h
       push 028858260h
       push 015124045h       
       push 05016A0C7h
       push 028191812h
       push 0F2401812h
       push 019154127h
       push 050F0F011h
       mov  ecx, 015124710h
       push ecx
       push 011151247h
       push 010111512h
       push 047101115h
       mov  eax, 012472015h
       push eax
       push eax
       push 012471A10h
       add  cl, 10h
       push ecx
       sub  cl, 20h
       push ecx

       xor  ecx, ecx
       dec  ecx

       // code starts
ps:  inc  ecx
       mov  edi, esp
go:  lodsb
       mov  bh, al
ft:  mov  ah, [edi]
       inc  edi
       shr  ah, 4
       sub  al, ah
       jnc  ft

       mov  al, [edi-1]
       and  al, 0Fh

       cmp  al, O_ERROR
       jnz  i7
       
       pop  edx
       not  edx

i7:  inc  edx
       cmp  al, O_UNIQUE
       jz   t_exit

       cmp  al, O_PREFIX
       jz   ps

       add  edi, 51h          //(_ettbl - _ttbl)

       cmp  al, O_EXTENDED
       jz   go

       mov  edi, [ebp+(8*4)+4]

i6:  inc  edx
       cmp  al, O_IMM8
       jz   t_exit
       cmp  al, O_MODRM
       jz   t_modrm
       cmp  al, O_WEIRD
       jz   t_weird

i5:  inc  edx
       cmp  al, O_IMM16
       jz   t_exit
       cmp  al, O_MODRM8
       jz   t_modrm

i4:  inc  edx
       cmp  al, O_IMM24
       jz   t_exit

i3:  inc  edx
i2:  inc  edx

       pushad
       mov  al, 66h
       repnz scasb
       popad
       jnz  c32

d2:  dec  edx
       dec  edx

c32: cmp  al, O_MODRM32
       jz   t_modrm
       sub  al, O_IMM32
       jz   t_imm32

i1:  inc  edx

t_exit:
       mov  esp, ebp
       mov  [esp+(7*4)], edx
       popad
       ret

//;*********************************
//;* PROCESS THE MOD/RM BYTE       *
//;*                               *
//;*   7    6 5          3 2    0  *
//;*   | MOD | Reg/Opcode | R/M |  *
//;*                               *
//;*********************************
t_modrm:
       lodsb
       mov  ah, al
       shr  al, 7
       jb   prmk
       jz   prm

       add  dl, 4

       pushad
       mov  al, 67h
       repnz scasb
       popad
       jnz  prm

d3:  sub  dl, 3

       dec  al
prmk:jnz  t_exit
       inc  edx
       inc  eax
prm:
       and  ah, 00000111b

       pushad
       mov  al, 67h
       repnz scasb
       popad
       jz   prm67chk

       cmp  ah, 04h
       jz   prmsib

       cmp  ah, 05h
       jnz  t_exit

prm5chk:
       dec  al
       jz   t_exit
i42: add  dl, 4
       jmp  t_exit

prm67chk:
       cmp  ax, 0600h
       jnz  t_exit
       inc  edx
       jmp  i1

prmsib:
       cmp  al, 00h
       jnz  i1
       lodsb
       and  al, 00000111b
       sub  al, 05h
       jnz  i1
       inc  edx
       jmp  i42

//;****************************
//;* PROCESS WEIRD OPCODES    *
//;*                          *
//;* Fucking test (F6h/F7h)   *
//;*                          *
//;****************************
t_weird:
       test byte ptr [esi], 00111000b
       jnz  t_modrm

       mov  al, O_MODRM8

       shr  bh, 1
       adc  al, 0
       jmp  i5

//;*********************************
//;* PROCESS SOME OTHER SHIT       *
//;*                               *
//;* Fucking mov (A0h/A1h/A2h/A3h) *
//;*                               *
//;*********************************
t_imm32:
       sub  bh, 0A0h

       cmp  bh, 04h
       jae  d2

       pushad
       mov  al, 67h
       repnz scasb
       popad
       jnz  chk66t

d4:  dec  edx
       dec  edx

chk66t:
       pushad
       mov  al, 66h
       repnz scasb
       popad
       jz   i1
       jnz  d2
}
}
unsigned int GetFunctionLength(__in void* begin)
{
	void *end = GetFunctionEnd(begin);
	unsigned int delta = (unsigned int)((DWORD_PTR)end - (DWORD_PTR)begin);
	delta    += mlde32(end);
	return delta;
}

void *GetFunctionEnd(__in void *func)
{
	void *block = func;
	TreeTrunk Tree = {0};
	Tree.MemTotalSize = 256;
	PVOID hHeap = RtlCreateHeap(HEAP_GROWABLE|HEAP_ZERO_MEMORY,0,Tree.MemTotalSize,Tree.MemTotalSize,0,0);
	Tree.MemStart = RtlAllocateHeap(hHeap,0,Tree.MemTotalSize);
	if(Tree.MemStart != NULL)
	{
		Tree.MemTotalSize = RtlSizeHeap(hHeap,0,Tree.MemStart);
		if(Tree.MemTotalSize != -1)
		{
			Tree.NumTotalEntry = Tree.MemTotalSize/4; 
			Tree.NumEntry = 0;
			//	ptr now points to the end of this block
			void *silenttree = GetBranchListFromBlock(block, &Tree);
			if(Tree.NumEntry == 0) return silenttree;
			DWORD Prev = NULL;
			for(DWORD i = 0; i < Tree.NumEntry; i++)
			{   
				DWORD sz = i * sizeof(DWORD);
				DWORD Limb = (*(DWORD*)Tree.MemStart + sz); 
				//	Skip branches that jump into a block we've already
				//	processed.  
				if(Limb < *(DWORD*)silenttree || Limb == Prev)
					continue;
				silenttree = GetFunctionEnd(&Limb);
				Prev     = Limb;
			}
			return silenttree;
		}
		RtlFreeHeap(hHeap,0,Tree.MemStart);
		return NULL;
	}
	return NULL;
}
void *GetBranchListFromBlock(__in void *block,__inout TreeTrunk *Tree)
{
	UCHAR *ptr = (UCHAR *)block;
	//	If we reach an end-point, then this block is complete
	while(!IsEndPoint(ptr, block))
	{
		//	Record all branching instructions that we encounter
		void *address = GetBranchAddress(ptr);
		if(address)
		{			
			*((DWORD*)Tree->MemStart + (Tree->NumEntry * 4)) = *(DWORD*)address;
			Tree->NumEntry++;	 
		}
		//	Next instruction
		ptr += mlde32(ptr);
	}
	return ptr;
}


void *GetBranchAddress(__in UCHAR *instr)
{
	long offset = 0;
	//	This code will determine what type of branch it is, and
	//	determine the address it will branch to.
	switch(*instr)
	{
	case INSTR_SHORTJMP:
	case INSTR_RELJCX:
		offset  = (long)(*(char *)(instr + 1));
		offset += 2;	
		break;
	case INSTR_RELJMP:
		offset  = *(long *)(instr + 1);
		offset += 5;
		break;
	case INSTR_NEAR_PREFIX:
		if(*(instr + 1) >= INSTR_NEARJCC_BEGIN && *(instr + 1) <= INSTR_NEARJCC_END)
		{
			offset  = *(long *)(instr + 2);
			offset += 5;
		}
		break;
	default:
		//	Check to see if it's in the valid range of JCC values.
		//	e.g. ja, je, jne, jb, etc..
		if(*instr >= INSTR_SHORTJCC_BEGIN && *instr <= INSTR_SHORTJCC_END)
		{
			offset  = (long)*((char *)(instr + 1));
			offset += 2;
		}
		break;
	}

	if(offset == 0) return NULL;
	return instr + offset;
}

__checkReturn bool IsEndPoint(__in UCHAR *instr,__in void *curblock)
{
	void *address;
	long offset;
	switch(*instr)
	{
	case INSTR_RET:
	case INSTR_RETN:
	case INSTR_RETFN:
	case INSTR_RETF:
		return true;
		break;

	//	The following two checks, look for an instance in which
	//	an unconditional jump returns us to a previous block,
	//	thus creating a pseudo-endpoint.
	case INSTR_SHORTJMP:
		offset  = (long)(*(char *)(instr + 1));
		address = instr + offset;
		if(address <= curblock) return true;
		break;
	case INSTR_RELJMP:
		offset  = *(long *)(instr + 1);
		address = instr + offset;
        if(address <= curblock) return true;
		break;
	default:
		return false;
		break;
	}

	return false;
}
PVOID MakeCodeCaveInDll(__in PUNICODE_STRING DllName,__in size_t bSize)
{
	PLDR_MODULE LdrInfo = {0};
	OBJECT_NAME_INFORMATION ObName = {0};
	SECTION_BASIC_INFORMATION sbInfo = {0};
	LARGE_INTEGER SecNewSize = {0};
	OBJECT_ATTRIBUTES oa = {0};
	HANDLE hTargetSection = INVALID_HANDLE_VALUE;
	NTSTATUS Status = 0;
	HMODULE hMod = 0;
	ULONG Result = 0;
	__try
	{
		Status = LdrFindEntryForAddress(GetPebDll(DllName->Buffer),&LdrInfo);
		if(!NT_SUCCESS(Status))
		{
			__leave;
		}
		if(wcscmp((wchar_t*)LdrInfo->BaseDllName.Buffer,DllName->Buffer) == 0)
		{
			Status = NtQueryObject(LdrInfo->BaseAddress,ObjectNameInformation,&ObName,sizeof(OBJECT_NAME_INFORMATION),&Result);
			if(!NT_SUCCESS(Status))
			{
				__leave;
			}
			InitializeObjectAttributes(&oa,&ObName.Name,OBJ_OPENIF,0,0);
			Status = NtOpenSection(&hTargetSection,SECTION_EXTEND_SIZE,&oa);
			if(!NT_SUCCESS(Status))
			{
				__leave;
			}
			Status = NtQuerySection(hTargetSection,SectionBasicInformation,&sbInfo,sizeof(SECTION_BASIC_INFORMATION),&Result);
			if(!NT_SUCCESS(Status))
			{
				__leave;
			}
			SecNewSize.LowPart =(sbInfo.Size.LowPart + (DWORD)bSize);
			Status = NtExtendSection(hTargetSection,&SecNewSize);
			if(!NT_SUCCESS(Status))
			{
				__leave;
			}
			return (PVOID)((ULONG)hMod + sbInfo.Size.LowPart);
		}
	}
	__except(1)
	{
		if(hTargetSection != INVALID_HANDLE_VALUE)
			NtClose(hTargetSection);
	}
	return 0;
}